import torch
import time
import cv2
import numpy as np
from PIL import Image

# 加载模型
print("正在加载模型...")
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')


# 获取类别名称
class_names = model.names

# 方法1：使用随机生成的张量
dummy_input = torch.randn(1, 3, 640, 640)
print(f"随机输入张量形状: {dummy_input.shape}")

# 计时并进行推理
start_time = time.time()
out_dummy = model(dummy_input)
dummy_inference_time = time.time() - start_time
print(f"随机输入推理时间: {dummy_inference_time:.4f} 秒")
print(f"输出结果形状: {[x.shape for x in out_dummy if hasattr(x, 'shape')]}")

# 方法2：使用真实图片
# 请替换为您的图片路径
image_path = r"C:\\Users\\masaikk\\Downloads\\zidane.jpg"  # 修改为你的实际图片路径

# 计时并进行推理
start_time = time.time()
# 直接将图片路径传给模型，不需要手动预处理
results = model(image_path)
inference_time = time.time() - start_time

print(f"实际图片推理时间: {inference_time:.4f} 秒")

# 获取原始图像用于绘制结果
img = cv2.imread(image_path)
if img is None:
    raise Exception("图片读取失败")

# 创建一个结果图像的副本用于绘制
result_img = img.copy()

# 获取图像原始尺寸
original_h, original_w = img.shape[:2]
print(f"原始图片形状: {img.shape}")

# 使用YOLOv5结果的pandas方法获取检测结果
detections = results.pandas().xyxy[0]  # 第一张图片的检测结果
print(f"检测到 {len(detections)} 个对象")

# 在原始图像上绘制检测框
for _, det in detections.iterrows():
    # 只处理置信度高于阈值的检测
    if det['confidence'] > 0.25:
        # 获取边界框坐标
        x1, y1, x2, y2 = int(det['xmin']), int(det['ymin']), int(det['xmax']), int(det['ymax'])
        
        # 获取置信度和类别
        conf = det['confidence']
        cls_id = det['class']
        name = det['name']
        
        # 绘制矩形框
        color = (0, 255, 0)  # 绿色框
        cv2.rectangle(result_img, (x1, y1), (x2, y2), color, 2)
        
        # 添加标签
        label = f"{name} {conf:.2f}"
        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(result_img, (x1, y1 - label_height - 10), (x1 + label_width, y1), color, -1)
        cv2.putText(result_img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

# 保存结果图像
cv2.imwrite('result.jpg', result_img)
print("检测结果已保存到 'result.jpg'")

# 显示结果图像（取消注释以显示）
# cv2.imshow('YOLOv5 Detection Result', result_img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

# 打印检测结果信息，方便调试
print("检测结果前5个:")
print(detections.head())

# 取消注释以下行可以导出为ONNX模型
# torch.onnx.export(model, dummy_input, 'yolov5n.onnx', opset_version=12)